import torch
from torch import nn
from torchvision import transforms
from torchsummary import summary
from dataset import GetDataLoader
import os
"""
卷积：卷积核=3x3，步幅=1，padding=1
池化：池化核=2x2，步幅=2
全连接层之间：一个激活函数，3个全连接，2个relu
VGG 使用的图片的大小是：224x224
"""

class VGG16(nn.Module):
    # 图片通道
    IMG_CHANNELS = GetDataLoader().IMG_CHANNELS
    # 图片分类数
    IMG_CLASSES = len(GetDataLoader().IMG_CLASSES)
    # 模型参数文件保存位置
    PTH_PATH = os.path.join('./model_files', 'best_model.pth')
    def __init_params(self):
        ## 初始化参数
        for m in self.modules():
            # 卷积层参数的初始化，w和b参数
            if isinstance(m, nn.Conv2d):
                ## 对于卷积，使用凯明初始化（对w初始化）
                nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
                ## 对b参数进行初始
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            # 全连接成参数的初始化
            elif isinstance(m, nn.Linear):
                ## 默认使用正态分布
                nn.init.normal_(m.weight, 0, 0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def __init__(self):
        super(VGG16,self).__init__()

        # 输入的图片是224x224x1
        ## 第一块，2卷积，2relu，1个maxpool
        self.block1 = nn.Sequential(
            nn.Conv2d(in_channels=self.IMG_CHANNELS,out_channels=64,kernel_size=3,padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=64,out_channels=64,kernel_size=3,padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,stride=2)
        )
        ## 第二块，2卷积，2relu，1个maxpool
        self.block2 = nn.Sequential(
            nn.Conv2d(in_channels=64,out_channels=128,kernel_size=3,padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=128,out_channels=128,kernel_size=3,padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,stride=2)
        )

        ## 第三块，3卷积，relu，1个maxpool
        self.block3 = nn.Sequential(
            nn.Conv2d(in_channels=128,out_channels=256,kernel_size=3,padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=256,out_channels=256,kernel_size=3,padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=256,out_channels=256,kernel_size=3,padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,stride=2)
        )

        ## 第四块，3卷积，relu，1个maxpool
        self.block4 = nn.Sequential(
            nn.Conv2d(in_channels=256,out_channels=512,kernel_size=3,padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=512,out_channels=512,kernel_size=3,padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=512,out_channels=512,kernel_size=2,padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,stride=2)
        )

        ## 第五块，3卷积，relu，1个maxpool
        self.block5 = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=2, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        ## 最后一块：1个展平，3个全连接
        self.block6 = nn.Sequential(
            nn.Flatten(),
            nn.Linear(in_features=7*7*512,out_features=4096),
            nn.ReLU(),
            nn.Linear(in_features=4096,out_features=4096),
            nn.ReLU(),
            nn.Linear(in_features=4096,out_features=self.IMG_CLASSES)
        )

        ## 初始化参数
        self.__init_params()

    def forward(self,input):
        input = self.block1(input)
        input = self.block2(input)
        input = self.block3(input)
        input = self.block4(input)
        input = self.block5(input)
        output = self.block6(input)

        return output


if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = VGG16()
    model = model.to(device)
    print(summary(model, (3, 224, 224)))